# 导入函数库
import numpy as np
import pandas as pd
import math
import statsmodels.api as sm
from jqdata import *
from dateutil.parser import parse
from typing import (Tuple,List,Dict,Union,Callable)
from tqdm import tqdm_notebook
from BuildPeriodDate import (GetTradePeriod,tdaysoffset)

def winsorize(factor, scale=3, have_negative=True):
    '''
    去极值函数
    factor:以股票code为index，因子值为value的Series
    std为几倍的标准差，have_negative 为布尔值，是否包括负值
    输出Series
    '''
    if isinstance(factor, pd.Series):
        r = factor.dropna().copy()
        if have_negative == False:
            r = r[r >= 0]
        else:
            pass
        # 取极值
        edge_up = r.mean() + scale * r.std()
        edge_low = r.mean() - scale * r.std()
        r[r > edge_up] = edge_up
        r[r < edge_low] = edge_low

        return r
    elif isinstance(factor, pd.DataFrame):
        r = factor.dropna().copy()
        if have_negative == False:
            r = r[r >= 0]
        else:
            pass
        # 取极值
        for col_index, col_value in r.iteritems():
            edge_up = col_value.mean() + scale * col_value.std()
            edge_low = col_value.mean() - scale * col_value.std()
            col_value[col_value > edge_up] = edge_up
            col_value[col_value < edge_low] = edge_low
            r[col_index] = col_value

        return r


# 中性化函数
# 传入：mkt_cap：以股票为index，市值为value的Series,
# factor：以股票code为index，因子值为value的Series,
# 输出：中性化后的因子值series
def neutralize(factor: pd.Series, mkt_cap=False, industry=True):
    y = factor
    if type(mkt_cap) == pd.Series:
        LnMktCap = mkt_cap.apply(lambda x: math.log(x))
        if industry:  # 行业、市值
            dummy_industry = getIndustryExposure(factor.index)
            x = pd.concat([LnMktCap, dummy_industry.T], axis=1)
        else:  # 仅市值
            x = LnMktCap
    elif industry:  # 仅行业
        dummy_industry = getIndustryExposure(factor.index)
        x = dummy_industry.T
    result = sm.OLS(y.astype(float), x.astype(float)).fit()
    return result.resid


# 为股票池添加行业标记,return df格式 ,为中性化函数的子函数
def getIndustryExposure(stock_list):
    df = pd.DataFrame(index=get_industries(name='sw_l1').index, columns=stock_list)
    for stock in stock_list:
        try:
            ndy_dict = get_industry(stock)
            ndy_cd = ndy_dict[stock]['sw_l1']['industry_code']
            df[stock][ndy_cd] = 1
        except:
            continue
    return df.fillna(0)  # 将NaN赋为0


'''过滤ST股票'''
def del_st_stock(securities: list, watch_date: str) -> list:
    info_ser = get_extras('is_st', securities,
                          end_date=watch_date, df=True, count=1).iloc[0]

    return info_ser[info_ser == False].dropna().index.tolist()


'''返回上市大于N日的股票'''
def del_iponday(securities: list, watch_date: str, N: int = 60) -> list:
    dic = {x: get_security_info(x, date=watch_date).start_date for x in securities if get_security_info(
        x, date=watch_date) != None}
    return [code for code, start_date in dic.items() if (parse(watch_date).date() - start_date).days > N]


'''过滤停牌股,默认为过滤当日停牌的'''
def del_pause(securities: list, watch_date: str, N: int = 10) -> list:
    paused = get_price(securities, end_date=watch_date, count=N, fields='paused', panel=False)
    if N == 1:
        return paused.query('paused != 1')['code'].values.tolist()
    else:

        paused_num = paused.groupby('code')['paused'].sum()
        return paused_num[paused_num < N].index.values.tolist()


'''获取股票池'''
def get_stockpool(symbol: str, watch_date: str) -> list:
    if symbol == 'A':
        stockList = get_index_stocks('000001.XSHG', date=watch_date) + get_index_stocks(
            '399107.XSHE', date=watch_date)
    else:
        stockList = get_index_stocks(symbol, date=watch_date)

    stockList = del_st_stock(stockList, watch_date)  # 过滤ST
    stockList = del_iponday(stockList, watch_date)  # 过滤上市不足60日
    stockList = del_pause(stockList, watch_date)  # 过滤当日停牌的股票
    return stockList


'''
因子计算
symbol:股票池所属成份，A为全市场
    start,end:回测区间 频率为月度
    last_date:计算next_returns的最后一日的日期
    factor:构建的因子
'''
def prepare_data(symbol: str, start: str, end: str, last_date: str, factor: Callable,
                 grid_func_params: Dict) -> pd.DataFrame:

    periods = GetTradePeriod(start, end, 'ME')  # 月度_周期获取
    periods = periods + [last_date]
    factor_dic = {}

    for tradeDate, nextDate in tqdm_notebook(zip(periods[:-1], periods[1:]), total=len(periods) - 1, desc='因子计算'):
        # 股票池
        stock_pool = get_stockpool(symbol, tradeDate.strftime('%Y-%m-%d'))

        # 因子获取

        f = factor(stock_pool, tradeDate, 20)

        # 获取数据 这里仅需获取一次就能使用该数据进行N此不同参数运算
        f.get_data()
        # 可以使用不同参数进行计算
        factor_frame = grid_func_params['func'](
            f, **grid_func_params['params'])

        # 计算下期收益率
        close = get_price(stock_pool, end_date=tradeDate, count=1,
                          fields='close', panel=False).set_index('code')

        next_close = get_price(stock_pool, end_date=nextDate,
                               count=1, fields='close', panel=False).set_index('code')

        next_returns = next_close['close'] / close['close'] - 1
        next_returns.name = 'next_ret'

        # 合并
        tmp = pd.concat((factor_frame, next_returns), axis=1)
        factor_dic[tradeDate] = tmp  # 放入容器

    factor_df = pd.concat(factor_dic)

    factor_df.index.names = ['date', 'code']

    return factor_df